// Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
// SPDX-License-Identifier: MIT

import { nanoid } from "nanoid";
import { toast } from "sonner";
import { create } from "zustand";
import { useShallow } from "zustand/react/shallow";

import { createClient } from "~/lib/supabase/client";

import { chatStream } from "../api";
import type { Message, MessageAgent, MessageRole, Resource } from "../messages";
import { mergeMessage } from "../messages";

import { getChatStreamSettings } from "./settings-store";

const createDefault = (threadId?: string) => ({
  loading: false,
  responding: false,
  threadId: threadId ?? nanoid(),
  messageIds: [],
  messages: new Map<string, Message>(),
  researcherMessageIds: [],
  firstReporterMessageId: null,
  resumeMatcherBeginMessageId: null,
  lastResumeMatcherKeyMessageId: null,
  resumeMatcherMessageIds: new Map<string, string[]>(),
  questionGeneratorBeginMessageId: null,
  lastQuestionGeneratorKeyMessageId: null,
  questionGeneratorMessageIds: new Map<string, string[]>(),
  interviewEvaluatorBeginMessageId: null,
  lastInterviewEvaluatorKeyMessageId: null,
  interviewEvaluatorMessageIds: new Map<string, string[]>(),
  ongoingMessageId: null,
  interviewPanelIsOpen: false,
  talentPoolDrawerIsOpen: false,
  jdGeneratorAcceptedMessageId: null,
  resumeMatcherAcceptedMessageId: null,
  questionGeneratorAcceptedMessageId: null,
});

export const useStore = create<{
  loading: boolean;
  responding: boolean;
  threadId: string;
  messageIds: string[];
  messages: Map<string, Message>;
  researcherMessageIds: string[];
  firstReporterMessageId: string | null;
  // resume matcher
  resumeMatcherBeginMessageId: string | null;
  lastResumeMatcherKeyMessageId: string | null;
  resumeMatcherMessageIds: Map<string, string[]>;
  // question generator
  questionGeneratorBeginMessageId: string | null;
  lastQuestionGeneratorKeyMessageId: string | null;
  questionGeneratorMessageIds: Map<string, string[]>;
  // interview evaluator
  interviewEvaluatorBeginMessageId: string | null;
  lastInterviewEvaluatorKeyMessageId: string | null;
  interviewEvaluatorMessageIds: Map<string, string[]>;
  ongoingMessageId: string | null;
  interviewPanelIsOpen: boolean;
  talentPoolDrawerIsOpen: boolean;

  jdGeneratorAcceptedMessageId: string | null;
  resumeMatcherAcceptedMessageId: string | null;
  questionGeneratorAcceptedMessageId: string | null;

  appendMessage: (message: Message) => void;
  updateMessage: (message: Message) => void;
  updateMessages: (messages: Message[]) => void;
  setOngoingMessage: (messageId: string | null) => void;
  openInterviewPanel: () => void;
  closeInterviewPanel: () => void;
  openTalentPoolDrawer: () => void;
  closeTalentPoolDrawer: () => void;
}>((set) => ({
  ...createDefault(),

  appendMessage(message: Message) {
    set((state) => ({
      messageIds: [...state.messageIds, message.id],
      messages: new Map(state.messages).set(message.id, message),
    }));
  },
  updateMessage(message: Message) {
    set((state) => ({
      messages: new Map(state.messages).set(message.id, message),
    }));
  },
  updateMessages(messages: Message[]) {
    set((state) => {
      const newMessages = new Map(state.messages);
      messages.forEach((m) => newMessages.set(m.id, m));
      return { messages: newMessages };
    });
  },
  setOngoingMessage(messageId: string | null) {
    set({ ongoingMessageId: messageId });
  },
  openInterviewPanel() {
    set({ interviewPanelIsOpen: true });
  },
  closeInterviewPanel() {
    set({ interviewPanelIsOpen: false });
  },
  openTalentPoolDrawer() {
    set({ talentPoolDrawerIsOpen: true });
  },
  closeTalentPoolDrawer() {
    set({ talentPoolDrawerIsOpen: false });
  },
}));

export async function sendMessage(
  content?: string,
  agentCode?: string,
  {
    role = "user",
    interruptFeedback,
    resources,
  }: {
    role?: MessageRole;
    interruptFeedback?: string;
    resources?: Array<Resource>;
  } = {},
  options: { abortSignal?: AbortSignal } = {},
) {
  if (content != null) {
    appendMessage({
      id: nanoid(),
      threadId: useStore.getState().threadId,
      role,
      content: content,
      contentChunks: [content],
      resources,
    });
  }

  const settings = getChatStreamSettings();
  const stream = chatStream(
    content ?? "[REPLAY]",
    {
      thread_id: useStore.getState().threadId,
      interrupt_feedback: interruptFeedback,
      resources,
      auto_accepted_plan: settings.autoAcceptedPlan,
      enable_background_investigation:
        settings.enableBackgroundInvestigation ?? true,
      max_plan_iterations: settings.maxPlanIterations,
      max_step_num: settings.maxStepNum,
      max_search_results: settings.maxSearchResults,
      mcp_settings: settings.mcpSettings,
      agent_code: agentCode,
    },
    options,
  );

  setResponding(true);
  let messageId: string | undefined;
  try {
    for await (const event of stream) {
      const { type, data } = event;
      messageId = data.id;
      let message: Message | undefined;
      if (type === "tool_call_result") {
        message = findMessageByToolCallId(data.tool_call_id);
      } else if (!existsMessage(messageId)) {
        message = {
          id: messageId,
          threadId: data.thread_id,
          agent: data.agent,
          role: data.role,
          content: "",
          contentChunks: [],
          isStreaming: true,
          interruptFeedback,
        };
        appendMessage(message);
      }
      message ??= getMessage(messageId);
      if (message) {
        message = mergeMessage(message, event);
        updateMessage(message);
      }
    }
  } catch {
    toast("An error occurred while generating the response. Please try again.");
    // Update message status.
    // TODO: const isAborted = (error as Error).name === "AbortError";
    if (messageId != null) {
      const message = getMessage(messageId);
      if (message?.isStreaming) {
        message.isStreaming = false;
        useStore.getState().updateMessage(message);
      }
    }
    useStore.getState().setOngoingMessage(null);
  } finally {
    setResponding(false);
  }
}

function setResponding(value: boolean) {
  useStore.setState({ responding: value });
}

function existsMessage(id: string) {
  return useStore.getState().messageIds.includes(id);
}

function getMessage(id: string) {
  return useStore.getState().messages.get(id);
}

function findMessageByToolCallId(toolCallId: string) {
  return Array.from(useStore.getState().messages.values())
    .reverse()
    .find((message) => {
      if (message.toolCalls) {
        return message.toolCalls.some((toolCall) => toolCall.id === toolCallId);
      }
      return false;
    });
}

function appendMessage(message: Message) {
  // 如果第一轮 research 没有生成 report，将 firstReporterMessageId 设置成 none
  if (message.finishReason === "interrupt") {
    if (
      useStore.getState().researcherMessageIds.length > 0 &&
      !useStore.getState().firstReporterMessageId
    ) {
      useStore.setState({
        firstReporterMessageId: "none",
      });
    }
  }
  if (message.agent === "coder" || message.agent === "researcher") {
    const current = useStore.getState().researcherMessageIds;
    useStore.setState({
      researcherMessageIds: [...current, message.id],
    });
    openInterviewPanel();
  } else if (
    message.agent === "reporter" &&
    !useStore.getState().firstReporterMessageId
  ) {
    useStore.setState({
      firstReporterMessageId: message.id,
    });
  } else if (
    message.agent === "resume_matcher" ||
    message.agent === "resume_matcher_summarize"
  ) {
    let resumeMatcherBeginMessageId =
      useStore.getState().resumeMatcherBeginMessageId;
    if (resumeMatcherBeginMessageId) {
      useStore.setState({
        resumeMatcherMessageIds: new Map(
          useStore.getState().resumeMatcherMessageIds,
        ).set(resumeMatcherBeginMessageId, [
          ...useStore
            .getState()
            .resumeMatcherMessageIds.get(resumeMatcherBeginMessageId)!,
          message.id,
        ]),
      });
    } else if (message.agent === "resume_matcher") {
      resumeMatcherBeginMessageId = message.id;
      useStore.setState({
        resumeMatcherBeginMessageId,
        lastResumeMatcherKeyMessageId: message.id,
        resumeMatcherMessageIds: new Map(
          useStore.getState().resumeMatcherMessageIds,
        ).set(resumeMatcherBeginMessageId, [message.id]),
      });
    }
  } else if (message.agent === "question_generator" && !message.toolCalls) {
    let questionGeneratorBeginMessageId =
      useStore.getState().questionGeneratorBeginMessageId;
    if (questionGeneratorBeginMessageId) {
      useStore.setState({
        questionGeneratorMessageIds: new Map(
          useStore.getState().questionGeneratorMessageIds,
        ).set(questionGeneratorBeginMessageId, [
          ...useStore
            .getState()
            .questionGeneratorMessageIds.get(questionGeneratorBeginMessageId)!,
          message.id,
        ]),
      });
    } else {
      questionGeneratorBeginMessageId = message.id;
      useStore.setState({
        questionGeneratorBeginMessageId,
        lastQuestionGeneratorKeyMessageId: message.id,
        questionGeneratorMessageIds: new Map(
          useStore.getState().questionGeneratorMessageIds,
        ).set(message.id, [message.id]),
      });
    }
  } else if (
    message.agent === "interview_evaluator" ||
    message.agent === "interview_evaluator_summarize"
  ) {
    let interviewEvaluatorBeginMessageId =
      useStore.getState().interviewEvaluatorBeginMessageId;
    if (interviewEvaluatorBeginMessageId) {
      useStore.setState({
        interviewEvaluatorMessageIds: new Map(
          useStore.getState().interviewEvaluatorMessageIds,
        ).set(interviewEvaluatorBeginMessageId, [
          ...useStore
            .getState()
            .interviewEvaluatorMessageIds.get(
              interviewEvaluatorBeginMessageId,
            )!,
          message.id,
        ]),
      });
    } else if (message.agent === "interview_evaluator") {
      interviewEvaluatorBeginMessageId = message.id;
      useStore.setState({
        interviewEvaluatorBeginMessageId,
        lastInterviewEvaluatorKeyMessageId: message.id,
        interviewEvaluatorMessageIds: new Map(
          useStore.getState().interviewEvaluatorMessageIds,
        ).set(message.id, [message.id]),
      });
    }
  }

  if (
    useStore.getState().resumeMatcherBeginMessageId &&
    message.agent !== "resume_matcher" &&
    message.agent !== "resume_matcher_summarize"
  ) {
    useStore.setState({
      resumeMatcherBeginMessageId: null,
    });
  }
  if (
    useStore.getState().questionGeneratorBeginMessageId &&
    message.agent !== "question_generator"
  ) {
    useStore.setState({
      questionGeneratorBeginMessageId: null,
    });
  }
  if (
    useStore.getState().interviewEvaluatorBeginMessageId &&
    message.agent !== "interview_evaluator" &&
    message.agent !== "interview_evaluator_summarize"
  ) {
    useStore.setState({
      interviewEvaluatorBeginMessageId: null,
    });
  }

  useStore.getState().setOngoingMessage(message.id);
  useStore.getState().appendMessage(message);
}

function updateMessage(message: Message) {
  if (getOngoingMessageId() && !message.isStreaming) {
    useStore.getState().setOngoingMessage(null);
  }
  useStore.getState().updateMessage(message);
}

function getOngoingMessageId() {
  return useStore.getState().ongoingMessageId;
}

export function openInterviewPanel() {
  useStore.getState().openInterviewPanel();
}

export function closeInterviewPanel() {
  useStore.getState().closeInterviewPanel();
}

export function openTalentPoolDrawer() {
  useStore.getState().openTalentPoolDrawer();
}

export function closeTalentPoolDrawer() {
  useStore.getState().closeTalentPoolDrawer();
}

export function useMessage(messageId: string | null | undefined) {
  return useStore(
    useShallow((state) =>
      messageId ? state.messages.get(messageId) : undefined,
    ),
  );
}

export function useMessageIds() {
  return useStore(useShallow((state) => state.messageIds));
}

export function useLastInterruptMessage() {
  return useStore(
    useShallow((state) => {
      if (state.messageIds.length >= 2) {
        const lastMessage = state.messages.get(
          state.messageIds[state.messageIds.length - 1]!,
        );
        return lastMessage?.finishReason === "interrupt" ? lastMessage : null;
      }
      return null;
    }),
  );
}

export function useLastFeedbackMessageIds() {
  const waitingForFeedbackMessageIds = useStore(
    useShallow((state) => {
      if (state.messageIds.length >= 2) {
        const lastMessage = state.messages.get(
          state.messageIds[state.messageIds.length - 1]!,
        );
        if (lastMessage && lastMessage.finishReason === "interrupt") {
          // Find the last interrupt message ID
          let prevInterruptIndex = -1;
          for (let i = state.messageIds.length - 2; i > 0; i--) {
            if (
              state.messages.get(state.messageIds[i]!)?.finishReason ===
              "interrupt"
            ) {
              prevInterruptIndex = i;
              break;
            }
          }

          // If no previous interrupt message found, return the second last message
          if (prevInterruptIndex === -1) {
            return [state.messageIds[state.messageIds.length - 2]];
          }

          // Get all run-- messages between last interrupt and current message
          const runMessageIds = state.messageIds
            .slice(prevInterruptIndex + 1, -1)
            .filter((id) => id.startsWith("run--"));

          return runMessageIds;
        }
      }
      return [];
    }),
  );
  return waitingForFeedbackMessageIds;
}

export function useLatestMessageId(agent: MessageAgent) {
  return useStore(
    useShallow((state) => {
      for (let i = state.messageIds.length - 1; i > 0; i--) {
        const messageId = state.messageIds[i]!;
        if (state.messages.get(messageId)?.agent === agent) {
          return messageId;
        }
      }
      return null;
    }),
  );
}

export function setJdGeneratorAccepted(messageId: string) {
  useStore.setState({
    jdGeneratorAcceptedMessageId: messageId,
  });
}

export function setResumeMatcherAccepted(messageId: string) {
  useStore.setState({
    resumeMatcherAcceptedMessageId: messageId,
  });
}

export function setQuestionGeneratorAccepted(messageId: string) {
  useStore.setState({
    questionGeneratorAcceptedMessageId: messageId,
  });
}

export function useToolCalls() {
  return useStore(
    useShallow((state) => {
      const toolCalls = state.messageIds
        ?.map((id) => getMessage(id)?.toolCalls)
        .filter((toolCalls) => toolCalls != null)
        .flat();
      return toolCalls ?? [];
    }),
  );
}

export async function retrieveMessages(agentCode: string, threadId?: string) {
  useStore.setState(createDefault(threadId));
  if (threadId) {
    useStore.setState({ loading: true });
    const supabase = createClient();
    const { data: pData } = await supabase
      .from("projects")
      .select("project_id")
      .eq("agent_code", agentCode)
      .eq("project_id", threadId)
      .eq("is_deleted", false);
    if (pData?.[0]?.project_id === threadId) {
      const { data } = await supabase
        .from("chat_messages")
        .select("*")
        .eq("project_id", threadId)
        .order("created_at", { ascending: true });
      data?.forEach((message) => {
        if (!existsMessage(message.id)) {
          const msg = {
            id: message.id,
            threadId,
            agent: message.agent,
            role: message.role,
            isStreaming: false,
            content: message.content,
            contentChunks: [],
            toolCalls: message.tool_calls,
            finishReason: message.finish_reason,
          };
          appendMessage(msg);
          updateMessage(msg);
        }
      });
    }
    useStore.setState({ loading: false });
  }
}
